package blsX

import (
	"crypto/sha256"

	bls12381 "github.com/kilic/bls12-381"
)

// 对不同一消息签名的聚合
// e(G, S) = e(G, s1+s2+s3 + ....)
// = e(G, k_1*h_1 + k_2*h_2 + k_3*h_3 ...)
// = e(G, k_1*h_1) * e(G, k_2*h_2) * e(G, k_3*h_3)
// = e(p1, h_1) * e(p2, h_2) * e(p3, h_3)
// S = s1 + s2 +s3
// s1 = k_i * h_i
// p_i = k_i * G  h_i = H(m_i)
// k_i 私钥， m_i 签名消息， p_i 公钥
func DiffMsgSignAggregate(msgs [][]byte, secretKeys [][32]byte) []byte {
	if len(msgs) != len(secretKeys) {
		panic("len of msgs is not equal len of secretKeys")
	}

	bls := bls12381.NewEngine()
	G2 := bls.G2
	res := G2.New()
	for i, m := range msgs {
		tmpG2, err := G2.HashToCurve(m, domainG2)
		if nil != err {
			panic(err)
		}
		var secret bls12381.Fr
		secret.FromBytes(secretKeys[i][:])

		G2.MulScalar(tmpG2, tmpG2, &secret)
		G2.Add(res, res, tmpG2)
	}

	return G2.ToCompressed(res)
}

func DiffMsgSignAggregateCheck(publicKeys [][48]byte, msgs [][]byte, sign [96]byte) bool {
	if len(msgs) != len(publicKeys) {
		panic("len of msgs is not equal len of publicKeys")
	}
	bls := bls12381.NewEngine()
	G2 := bls.G2
	G1 := bls.G1

	for i, m := range msgs {
		tmpG2, err := G2.HashToCurve(m, domainG2)
		if nil != err {
			panic(err)
		}
		tmpG1, err := G1.FromCompressed(publicKeys[i][:])
		if nil != err {
			panic(err)
		}

		bls.AddPair(tmpG1, tmpG2)
	}
	sigG2, err := G2.FromCompressed(sign[:])
	if nil != err {
		panic(err)
	}
	bls.AddPairInv(G1.One(), sigG2)
	return bls.Check()
}

// sign for msg
func SingleSign(msg []byte, secret [32]byte) []byte {
	// msg 映射到 G2
	bls := bls12381.NewEngine()
	G2 := bls.G2
	tmpG2, err := G2.HashToCurve(msg, domainG2)
	if nil != err {
		panic(err)
	}

	// secret to Fr
	var secretFr bls12381.Fr
	secretFr.FromBytes(secret[:])
	sign := G2.MulScalar(tmpG2, tmpG2, &secretFr)
	return G2.ToCompressed(sign)
}

// check sign
func SingleCheck(sign [96]byte, publickey [48]byte, msg []byte) bool {
	bls := bls12381.NewEngine()
	G2 := bls.G2
	G1 := bls.G1

	// sign to P2
	signP2, err := G2.FromCompressed(sign[:])
	if nil != err {
		panic(err)
	}

	bls.AddPair(G1.One(), signP2)
	// public to P1
	pubP1, err := G1.FromCompressed(publickey[:])
	if nil != err {
		panic(err)
	}

	// msg to P2
	msgP2, err := G2.HashToCurve(msg, domainG2)
	if nil != err {
		panic(err)
	}

	bls.AddPairInv(pubP1, msgP2)

	return bls.Check()
}

// 对同一消息签名的n-n聚合（即n-n多重签名)
func MsgSignAggregateNN(secretKeys [][32]byte, msg []byte) []byte {
	bls := bls12381.NewEngine()
	G2 := bls.G2

	tmpG2, err := G2.HashToCurve(msg, domainG2)
	if nil != err {
		panic(err)
	}

	n := len(secretKeys)
	pubkey := make([][]byte, n)
	superp := make([]byte, n*48)
	secretFr := make([]bls12381.Fr, n)
	fr := bls12381.NewFr()
	for i := 0; i < n; i++ {
		secretFr[i] = *fr.FromBytes(secretKeys[i][:])
		pubkey[i] = GetPublicKey(secretKeys[i])
		copy(superp[48*i:48*(i+1)], secretKeys[i][:])
	}
	hash := sha256.New()
	poly := make([]bls12381.Fr, len(secretKeys))
	// ai = H(pi || {p1 || p2 || p3})
	for i := 0; i < n; i++ {
		hash.Write(pubkey[i])
		hash.Write(superp)
		var tmpByte []byte
		hash.Sum(tmpByte)
		hash.Reset()
		var tmpFr bls12381.Fr
		tmpFr.FromBytes(tmpByte)
		poly[i] = tmpFr
	}

	sumFr := fr.Zero()
	// sum (ai * pki)
	for i := 0; i < n; i++ {
		var res bls12381.Fr
		res.Mul(&poly[i], &secretFr[i])
		sumFr.Add(sumFr, &res)
	}
	//
	G2.MulScalar(tmpG2, tmpG2, sumFr)
	return G2.ToCompressed(tmpG2)
}

func MsgSignAggregateNNCheck(pubkeys [][48]byte, msg []byte, sign [96]byte) bool {
	bls := bls12381.NewEngine()
	G2 := bls.G2
	G1 := bls.G1

	n := len(pubkeys)
	hash := sha256.New()
	poly := make([]bls12381.Fr, len(pubkeys))

	superp := make([]byte, 0, n*48)
	for _, p := range pubkeys {
		superp = append(superp, p[:]...)
	}

	// ai = H(pi || {p1 || p2 || p3})
	for i := 0; i < n; i++ {
		hash.Write(pubkeys[i][:])
		hash.Write(superp)
		var tmpByte []byte
		hash.Sum(tmpByte)
		hash.Reset()
		var tmpFr bls12381.Fr
		tmpFr.FromBytes(tmpByte)
		poly[i] = tmpFr
	}

	sumG1 := G1.Zero()
	for i := 0; i < n; i++ {
		p, err := G1.FromCompressed(pubkeys[i][:])
		if nil != err {
			panic(err)
		}
		G1.MulScalar(p, p, &poly[i])
		G1.Add(sumG1, sumG1, p)
	}
	signP2, err := G2.FromCompressed(sign[:])
	if nil != err {
		panic(err)
	}

	// msg to P2
	msgP2, err := G2.HashToCurve(msg, domainG2)
	if nil != err {
		panic(err)
	}

	bls.AddPair(G1.Zero(), signP2)
	bls.AddPairInv(sumG1, msgP2)
	return bls.Check()
}
